#!/usr/bin/env python


__author__ = 'Ben "TheX1le" Smith'
__email__  = 'thex1le _A_T_ remote-exploit.org'
__website__= 'remote-exploit.org'
__date__   = '05/18/2011'
__version__= '2.0.2'
__file__   = 'airgraph-ng'
__data__   = 'This is the main airgraph-ng file'
import subprocess
import sys, argparse, os
import pdb

try:
    from airgraphviz import *
except ImportError as error:
    raise Exception("Your airgraph-ng installation is broken, could not import libraries: %s" %(error))

path = libOuiParse.path

class interface:
    """
    provide basic UI to user
    """
    def header(self):
        """
        Print a pretty header out
        """
        print('#'*42+"\n#"+" "*9+"Welcome to Airgraph-ng"+" "*9+"#\n"+"#"*42+"\n")

class dotCreate:
    """
    Class for creating graphviz .dot config files
    """
    def __init__(self, csv_, png_):
        ret = libDumpParse.airDumpParse().parser(csv_)
        self.capr   = ret['capr']
        self.Ap     = ret['apDict']
        self.client = ret['clientDict']
        self.NA     = ret['NA']
        self.NAP    = ret['NAP']
        self.csv    = csv_
        self.png    = png_

        #144x144 hard code image size to 12feet x 12feet
        #start graphviz config file
        self.dotFile = ['digraph G{\n\tsize ="144,144";\n\toverlap=false;\n']
        self.ouiCheck = libOuiParse.macOUI_lookup(path + 'oui.txt')
        if self.ouiCheck is False:
            print("Missing the oui.txt file from http://standards-oui.ieee.org/oui.txt, place it in the support directory")
            sys.exit(-1)
        self.ouiCheck.identDeviceDict(path + 'ouiDevice.txt')

    def CAPR(self):
        """
        Client AP relationship graph
        Display a graph showing what clients are talking to what AP's
        """
        tclient = 0 #total client number
        tap     = len(list(self.capr.keys())) #total ap list
        for bssid in list(self.capr.keys()):
            time = [self.Ap[bssid]['lts'],self.Ap[bssid]['fts']]
            priv = self.Ap[bssid]['privacy']
            if priv == '':
                priv = "Unknown"
            Color = self.encColor(priv)
            tclient += len(self.capr[bssid])
            for client in self.capr[bssid]:
                self.dotFile.append(self.linker(bssid,'->',client))
                self.dotFile.append(self.clientColor(client,'black',[self.client[client]['lts'],self.client[client]['fts']],None,True))
            self.dotFile.append(
                self.apLabel(bssid,Color,len(self.capr[bssid]),time))
        footer ='label="Generated by Airgraph-ng\\n%s Access Points and\\n%s Clients shown";\n}' %(tap,tclient)
        self.dotFile.append(footer)
        self.dot = os.path.splitext(self.csv)[0] + '-capr.dot'
        if self.png == None:
            self.png = os.path.splitext(self.csv)[0] + '-capr.png'


    def CPG(self):
        """
        Common Probe Graph
        Shows a graph of every client requesting similar probes
        """
        probeCount = 0
        clientCount = {}
        for key in list(self.client.keys()):
            sdata = self.client[key]
            lts = sdata['lts']
            fts = sdata['fts'] 
            if sdata["probe"] != ['']:
                lpc = len(sdata['probe'])
                clientCount[key] = key
                probeCount += lpc
                for probe in sdata['probe']:
                    cleaned_probe = probe.replace('"', '\\"')
                    self.dotFile.append(self.clientColor(cleaned_probe,'blue'))
                    self.dotFile.append(self.linker(key,'->', cleaned_probe))
                    clientColor = '%s\\nRequested %s probes' %(key,lpc)
                    self.dotFile.append(self.clientColor(key,'black',[lts,fts],None,True))
        footer = 'label="Generated by Airgraph-ng\\n%s Probes and \\n%s Clients are shown";\n}' % (probeCount,len(list(clientCount.keys())))
        self.dotFile.append(footer)
        self.dot = os.path.splitext(self.csv)[0] + '-cpg.dot'
        if self.png == None:
            self.png = os.path.splitext(self.csv)[0] + '-cpg.png'

    def clientColor(self,mac,color,time=None,label=None,DouiCheck=False):
        """
        format the client with a color and a label
        return graphiz format line
        """

        if label == None:
            label = mac
        #device OUI check
        if DouiCheck == True:
            lts = time[0]
            fts = time[1]
            rtn = '\t"%s" [label="%s\\nOUI: %s\\nDevice Type: %s\\nFirst Time Seen: %s\\nLast Time Seen: %s" color=%s fontsize=9];\n' %(mac,mac,self.APouiLookup(mac),self.clientOuiLookup(mac),fts,lts,color)
        else:
            rtn = '\t"%s" [label="%s" color=%s fontsize=9];\n' %(mac,label,color)

        return rtn

    def encColor(self,enc):
        """
        Take encryption type and decide what color it should be displayed as
        returns a list containing AP fill color and Ap label font color
        """
        fontColor = "black" #default font color

        if enc =="OPN":
            color = "firebrick2"
        elif enc == "WEP":
            color = "gold2"
        elif enc in ["WPA","WPA2WPA","WPA2","WPAOPN", "WPA3 WPA2"]:
            color = "green3"
        else: #unknown enc type
            color = "black"
            fontColor = "white"

        return (color,fontColor)

    def linker(self,objA,sep,objB):
        """
        Return a graphviz line that links 2 objects together
        Both objects are passed in with a separator
        returns graphiz format line
        """
        return '\t"%s"%s"%s";\n' %(objA,sep,objB)

    def dotWrite(self):
        """
        Write all the information obtained to a config file
        """
        dotdata = ''.join(self.dotFile)
        try:
            os.remove(self.dot)
        except Exception:
            pass
        nfile = open(self.dot,'a')
        nfile.write(dotdata)
        nfile.close()

    def clientOuiLookup(self,oui):
        """
        check ouiDevices and attempt to determine the device type
        """
        prefix = oui[:8]
        if prefix in self.ouiCheck.ouitodevice:
            device = self.ouiCheck.ouitodevice[prefix]
        else:
            device = 'Unknown'

        return device

    def APouiLookup(self,oui):
        """
        check the oui.txt file and determine the manufacture for an AP or client
        """
        prefix = oui[:8]
        company = self.ouiCheck.lookup_OUI(prefix)
        if company == False:
            company = 'Unknown'
        return company

    def apLabel(self,bssid,color,cnum,time):
        """
        Create label strings for AP's
        """
        lts = time[0]
        fts = time[1]
        AP = self.Ap[bssid]
        return'\t"%s" [label="%s\\nEssid: %s\\nChannel: %s\\nEncryption: %s\\nOUI: %s\\nFirst Time Seen: %s\\n Last Time Seen: %s\\nNumber of Clients: %s" style=filled fillcolor="%s" fontcolor="%s" fontsize=9];\n' %(bssid,bssid,AP['essid'].rstrip('\x00').replace('"', '\\"'),AP['channel'],AP['privacy'],self.APouiLookup(bssid),fts,lts,cnum,color[0],color[1])

    def graphCreate(self,keepdot_):
        """
        Write out the graphviz dotFile and creat the graph
        """
        print("\n**** WARNING Images can be large, up to 12 Feet by 12 Feet****")
        print("Creating your Graph using, %s and writing to, %s"  %(self.csv,self.png))
        print("Depending on your system this can take a bit. Please standby......")

        try:
            subprocess.Popen(["fdp","-Tpng",self.dot,"-o",self.png,"-Gcharset=latin1"]).wait()
        except Exception:
            os.remove(self.dot)
            print("You seem to be missing the Graphviz toolset. Did you check out the airgraph-ng Deps in the readme?")
            sys.exit(1)
        if not keepdot_:
            os.remove(self.dot)

if __name__ == "__main__":
    """
    Main function.
    Parses command line input for proper switches and arguments. 
    Error checking is done in here.
    Variables are defined and all calls are made from MAIN.
    """
    parser = argparse.ArgumentParser(description="Generate Client to AP Relationship (CAPR) and Common probe graph (CPG) from a airodump-ng CSV file")  #
    parser.add_argument("-o",
        "--output",
        help="Our Output Image ie... Image.png")
    parser.add_argument("-i",
        "--input",
        help="Airodump-ng txt file in CSV format. NOT the pcap")
    parser.add_argument("-g",
        "--graph",
        dest="graph_type",
        help="Graph Type Current [CAPR (Client to AP Relationship) OR CPG (Common probe graph)]")
    parser.add_argument("-d",
        "--dotfile",
        dest="keep_dot",
        action="store_true",
        default=False,
        help="Keep the dot graph file after the export to the PNG image has been done")

    if len(sys.argv) <= 1:
        interface().header()
        parser.print_help()
        sys.exit(0)

    args = parser.parse_args()

    outFile = args.output
    graphType = args.graph_type
    inFile = args.input
    keepdot = args.keep_dot

    if inFile == None:
        print("Error No Input File Specified")
        sys.exit(1)

    if graphType == None:
        print("Error No Graph Type Defined")
        sys.exit(1)

    if graphType.upper() not in ['CAPR','CPG','ZKS']:
        print("Error Invalid Graph Type\nVaild types are CAPR or CPG")
        sys.exit(1)

    dot = dotCreate(inFile, outFile)

    if graphType.upper() == 'CPG':
        dot.CPG()
    elif graphType.upper() == 'CAPR':
        dot.CAPR()

    dot.dotWrite()
    dot.graphCreate(keepdot)
